a5f94c
@@ -527,14 +527,15 @@
private OperationHandle executeStatementInternal(String statement,
       operation = getOperationManager().newExecuteStatementOperation(getSession(), statement,
           confOverlay, runAsync, queryTimeout);
       opHandle = operation.getHandle();
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
       // Refering to SQLOperation.java, there is no chance that a HiveSQLException throws and the
       // async background operation submits to thread pool successfully at the same time. So, Cleanup
       // opHandle directly when got HiveSQLException
       if (opHandle != null) {
+        removeOpHandle(opHandle);
         getOperationManager().closeOperation(opHandle);
       }
       throw e;
@@ -567,10 +568,11 @@
public OperationHandle getTypeInfo()
     GetTypeInfoOperation operation = operationManager.newGetTypeInfoOperation(getSession());
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -587,10 +589,11 @@
public OperationHandle getCatalogs()
     GetCatalogsOperation operation = operationManager.newGetCatalogsOperation(getSession());
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -608,10 +611,11 @@
public OperationHandle getSchemas(String catalogName, String schemaName)
         operationManager.newGetSchemasOperation(getSession(), catalogName, schemaName);
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -630,10 +634,11 @@
public OperationHandle getTables(String catalogName, String schemaName, String t
         operationManager.newGetTablesOperation(getSession(), catalogName, schemaName, tableName, tableTypes);
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -650,10 +655,11 @@
public OperationHandle getTableTypes()
     GetTableTypesOperation operation = operationManager.newGetTableTypesOperation(getSession());
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -675,10 +681,11 @@
public OperationHandle getColumns(String catalogName, String schemaName,
         catalogName, schemaName, tableName, columnName);
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -692,6 +699,12 @@
private void addOpHandle(OperationHandle opHandle) {
     }
   }
 
+  private void removeOpHandle(OperationHandle opHandle) {
+    synchronized (opHandleSet) {
+      opHandleSet.remove(opHandle);
+    }
+  }
+
   @Override
   public OperationHandle getFunctions(String catalogName, String schemaName, String functionName)
       throws HiveSQLException {
@@ -702,10 +715,11 @@
public OperationHandle getFunctions(String catalogName, String schemaName, Strin
         .newGetFunctionsOperation(getSession(), catalogName, schemaName, functionName);
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -940,10 +954,11 @@
public OperationHandle getPrimaryKeys(String catalog, String schema,
         .newGetPrimaryKeysOperation(getSession(), catalog, schema, table);
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
@@ -964,10 +979,11 @@
public OperationHandle getCrossReference(String primaryCatalog,
          foreignSchema, foreignTable);
     OperationHandle opHandle = operation.getHandle();
     try {
-      operation.run();
       addOpHandle(opHandle);
+      operation.run();
       return opHandle;
     } catch (HiveSQLException e) {
+      removeOpHandle(opHandle);
       operationManager.closeOperation(opHandle);
       throw e;
     } finally {
